from keras.layers import Input, Embedding, Conv1D, Dense, Dropout, GlobalMaxPooling1D, Concatenate, Multiply,Dot,Lambda
from keras.models import Model
#from tensorflow import keras as K
from tensorflow.contrib import keras as K
from keras import backend as back
from keras import optimizers
import numpy as np

def create_net(embedding_dim, char_size, num_filters,
               embeddings_matrix, hidden_dim, input_length,
               sub_mask, detail_mask,
               intent_output_dim = 4,
               subtask_output_dim = 10,
               subtask_detail_output_dim = 64):


    input_sentence = Input(shape=(input_length,), name='input_sentence')
    # sub_mask_input = Input(shape=(4, 10), name='sub_mask')
    # detail_mask_input = Input(shape=(10, 64), name='detail_mask')


    char_embeding = Embedding(output_dim=embedding_dim,
                              input_dim =char_size,
                              weights   =[embeddings_matrix],  # 重点：预训练的词向量
                              input_length=input_length,
                              name='char_embeding',
                              trainable=False)(input_sentence)



    char_conv_1 = Conv1D(filters=num_filters, kernel_size=1, activation=K.backend.relu,
                         name='char_conv_1')(char_embeding)
    char_conv_2 = Conv1D(filters=num_filters, kernel_size=2, activation=K.backend.relu,
                         name='char_conv_2')(char_embeding)
    char_conv_3 = Conv1D(filters=num_filters, kernel_size=3, activation=K.backend.relu,
                         name='word_conv_3')(char_embeding)
    char_conv_4 = Conv1D(filters=num_filters, kernel_size=4, activation=K.backend.relu,
                         name='char_conv_4')(char_embeding)
    char_conv_5 = Conv1D(filters=num_filters, kernel_size=5, activation=K.backend.relu,
                         name='char_conv_5')(char_embeding)

    char_out_1 = GlobalMaxPooling1D(name='char_pool_1')(char_conv_1)
    char_out_2 = GlobalMaxPooling1D(name='char_pool_2')(char_conv_2)
    char_out_3 = GlobalMaxPooling1D(name='char_pool_3')(char_conv_3)
    char_out_4 = GlobalMaxPooling1D(name='char_pool_4')(char_conv_4)
    char_out_5 = GlobalMaxPooling1D(name='char_pool_5')(char_conv_5)

    char_word = Concatenate(axis=1)([char_out_1,
                                     char_out_2,
                                     char_out_3,
                                     char_out_4,
                                     char_out_5
                                     ])

    #同时进行三个分类的输出
    #意图
    intent_output = Dense(units=hidden_dim, activation=K.backend.relu)(char_word)
    intent_output = Dropout(rate=0.5)(intent_output)
    intent_output = Dense(units=intent_output_dim, activation=K.backend.softmax,
                          name = "intent_output")(intent_output)

    # sub_mask = sub_mask.transpose()
    # #增加类别矩阵，用于规范类别
    sub_mask = back.constant(sub_mask)
    sub_mask_layer = Lambda(lambda x: back.dot(x,sub_mask))
    indent_mask = sub_mask_layer(intent_output)

    #子任务
    subtask_output = Dense(units=hidden_dim, activation=K.backend.relu)(char_word)
    subtask_output = Dropout(rate=0.5)(subtask_output)
    subtask_output = Dense(units=subtask_output_dim, activation=K.backend.softmax,
                           name="subtask_output_pre")(subtask_output)

    # indent_mask = back.transpose(indent_mask)
    subtask_output = Multiply(name = "subtask_output")([subtask_output, indent_mask])

    detail_mask = back.constant(detail_mask)
    detail_mask_layer = Lambda(lambda x: back.dot(x,detail_mask))

    detail_mask = detail_mask_layer(subtask_output)
    #子任务标签
    subtask_detail_output = Dense(units=hidden_dim, activation=K.backend.relu)(char_word)
    subtask_detail_output = Dropout(rate=0.5)(subtask_detail_output)
    subtask_detail_output = Dense(units=subtask_detail_output_dim, activation=K.backend.softmax,
                                  name="subtask_detail_output_pre")(subtask_detail_output)

    #进行转置 匹配维度
    # detail_mask = back.transpose(detail_mask)
    subtask_detail_output = Multiply(name = "subtask_detail_output")([subtask_detail_output, detail_mask])




    model = Model(inputs=[input_sentence,], outputs=[intent_output, subtask_output, subtask_detail_output])


    model.summary()
    nadam = optimizers.Nadam()

    #一个多任务学习的模型，同时学习三个分类器
    model.compile(optimizer=nadam,

                  loss={'intent_output': 'categorical_crossentropy',
                        'subtask_output': 'categorical_crossentropy',
                        'subtask_detail_output':'categorical_crossentropy'},

                  loss_weights={'intent_output':0.5,
                                'subtask_output':0.35,
                                'subtask_detail_output':0.15},

                  metrics=['accuracy'])
    return model
